//
// Created by hjie on 23-12-27.
//

#include "codec_context.h"
#include <libavutil/avutil.h>
#include <libavutil/samplefmt.h>
#include <sstream>
#include "base_util.h" 
#include "input_output_context.h"
#include "hw_instance.h"
#include "media_stream.h"

DecoderContext::DecoderContext()
{
    m_decoder_ctx = nullptr;
    m_codec_hw_config = nullptr;
    m_hw_frame_init_status = true;
}

DecoderContext::~DecoderContext()
{
    if (m_decoder_ctx)
    {
        avcodec_close(m_decoder_ctx);
        avcodec_free_context(&m_decoder_ctx);
        m_decoder_ctx = nullptr;
        PrintLog("release codec ctx resouce");
    }
}

bool DecoderContext::OpenCodecCtx()
{
    if (m_codec_hw_config)
    {
        m_hw_frame_init_status = false;
        bool ret = ProcessHwConfig();
        if (!ret)
        {
            return false;
        }
    }
    else
    {
        m_hw_frame_init_status = true;
    }
    printf("open codec before format : %d.\n", m_decoder_ctx->pix_fmt);
    int result = avcodec_open2(m_decoder_ctx, m_decoder_ctx->codec, nullptr);
    if (result < 0)
    {
        std::stringstream ss;
        ss << "ctx codec id : " << m_decoder_ctx->codec_id << "\tcodec id : " << m_decoder_ctx->codec->id;
        PrintLog(result);
        PrintLog(ss.str());
        return false;
    }
    printf("open codec after format : %d.\n", m_decoder_ctx->pix_fmt);
    return true;
}

bool DecoderContext::ProcessHwFrameStatus()
{
    if (m_hw_frane_callback)
    {
        return m_hw_frane_callback();
    }
    return false;
}

void DecoderContext::HwFrameFunc(std::function<bool()> &&callback_func)
{
    m_hw_frane_callback = callback_func;
}

bool DecoderContext::ProcessHwConfig()
{
    if (m_codec_hw_config && m_codec_hw_config->methods & AV_CODEC_HW_CONFIG_METHOD_HW_DEVICE_CTX)
    {
        m_decoder_ctx->hw_device_ctx = HwInstance::Instance()->GetBuffer(m_codec_hw_config->device_type);
        if (!m_decoder_ctx->hw_device_ctx)
        {
            return false;
        }
        m_decoder_ctx->opaque     = this;
        m_decoder_ctx->get_format = [](AVCodecContext * codec_ctx, const AVPixelFormat * pix_fmt) {

            DecoderContext * data = (DecoderContext*)codec_ctx->opaque;
            if (data->m_codec_hw_config->device_type == AV_HWDEVICE_TYPE_QSV)
            {
                return AV_PIX_FMT_NV12;
            }
            else if (data->m_codec_hw_config->device_type == AV_HWDEVICE_TYPE_VAAPI)
            {
                return AV_PIX_FMT_YUV420P;
            }
            return AV_PIX_FMT_NONE;
        };
        m_hw_frame_init_status = true;
    }
    else if (m_codec_hw_config && m_codec_hw_config->methods & AV_CODEC_HW_CONFIG_METHOD_HW_FRAMES_CTX)
    {
        m_hw_frame_init_status = false;
    }
    return true;
}

void DecoderContext::ConfigParam(std::shared_ptr<MediaStream> &media_stream)
{
    m_decoder_ctx->pkt_timebase = media_stream->AvStream()->time_base;
    if (media_stream->AvStream()->codecpar->codec_type == AVMEDIA_TYPE_AUDIO)
    {
        m_decoder_ctx->sample_fmt = (AVSampleFormat)media_stream->AvStream()->codecpar->format;
    }
    else
    { 
        m_decoder_ctx->pix_fmt      = (AVPixelFormat)media_stream->AvStream()->codecpar->format;
    }
    m_decoder_ctx->framerate    = media_stream->AvStream()->avg_frame_rate;
    m_media_stream = media_stream;
}

bool DecoderContext::CreateDecoderContextByName(const std::string &name, const std::string & hw_name)
{
    const AVCodec * codec = nullptr;
    if (!hw_name.empty())
    {
        codec = FindCodecByNameAndHw(name);
    }
    else
    {
        codec = FindCodecByName(name);
    }
    if (!codec)
    {
        return false;
    }
    m_decoder_ctx = avcodec_alloc_context3(codec);
    if (!m_decoder_ctx)
    {
        m_decoder_ctx = nullptr;
        return false;
    }
    return true;
}

bool DecoderContext::ConfigHw(const std::string &hw_name)
{
    m_decoder_ctx->opaque = (void*)this;
    m_decoder_ctx->get_format = [](AVCodecContext * codec_ctx, const AVPixelFormat * av_pixel_formats){

        auto * opaque = (DecoderContext*)codec_ctx->opaque;
        if (!opaque)
        {
            return AV_PIX_FMT_NONE;
        }
        const AVPixelFormat * p = nullptr;
        for (p = av_pixel_formats; *p != AV_PIX_FMT_NONE; p++)
        {
            const AVPixFmtDescriptor *desc = av_pix_fmt_desc_get(*p);
            const AVCodecHWConfig  *config = NULL;
            int i;

            if (!(desc->flags & AV_PIX_FMT_FLAG_HWACCEL))
            {
                break;
            }
            {
                for (i = 0;; i++)
                {
                    config = avcodec_get_hw_config(opaque->m_decoder_ctx->codec, i);
                    if (!config)
                        break;
                    if (!(config->methods &
                          AV_CODEC_HW_CONFIG_METHOD_HW_DEVICE_CTX))
                        continue;
                    if (config->pix_fmt == *p)
                        break;
                }
            }
            if (config && config->device_type == AV_HWDEVICE_TYPE_VAAPI)
            {
                opaque->m_codec_hw_config = config;
                // int ret = hwaccel_decode_init(s);
                // if (ret < 0) {
                //     if (ist->hwaccel_id == HWACCEL_GENERIC) {
                //         av_log(NULL, AV_LOG_FATAL,
                //                "%s hwaccel requested for input stream #%d:%d, "
                //                "but cannot be initialized.\n",
                //                av_hwdevice_get_type_name(config->device_type),
                //                ist->file_index, ist->st->index);
                //         return AV_PIX_FMT_NONE;
                //     }
                //     continue;
                // }
                // ist->hwaccel_pix_fmt = *p;
                break;
            }
        }
        return *p;
    };
    // if (m_codec_hw_config && m_codec_hw_config->methods & AV_CODEC_HW_CONFIG_METHOD_HW_DEVICE_CTX)
    {
        m_decoder_ctx->hw_device_ctx = HwInstance::Instance()->GetBuffer(hw_name);
    }
    return true;
}

bool DecoderContext::CreateDecoderContextById(AVCodecID codec_id, const std::string & hw_name)
{
    const AVCodec * codec = nullptr;
    if (!hw_name.empty())
    {
        codec = FindCodecByIdAndHw(codec_id, HwInstance::Instance()->GetHwType(hw_name));
    }
    else
    {
        codec = FindCodecById(codec_id);
    }
    if (!codec)
    {
        PrintLogMsg("no found codec");
        return false;
    }
    m_decoder_ctx = avcodec_alloc_context3(codec);
    if (!m_decoder_ctx)
    {
        m_decoder_ctx = nullptr;
        return false;
    }
    return true;
}


const AVCodec * DecoderContext::FindCodecById(AVCodecID codec_id)
{
    auto target = FindCodecByCustomFunc([codec_id](const AVCodec * codec){

        if (codec && codec->id == codec_id && (codec->decode != nullptr || codec->receive_frame != nullptr))
        {
            return true;
        }
        return false;
    });
    return target;
}


const AVCodec * DecoderContext::FindCodecByName(const std::string &name)
{

    const AVCodec * target = FindCodecByCustomFunc([name](const AVCodec * codec){

        if (codec && name == codec->name && (codec->decode != nullptr || codec->receive_frame != nullptr))
        {
            return true;
        }
        return false;

    });
    return target;
}

const AVCodec * DecoderContext::FindCodecByNameAndHw(const std::string & name, const std::string & hw_name)
{
    AVHWDeviceType hw_device_type = AV_HWDEVICE_TYPE_NONE;
    if (!hw_name.empty())
    {
        hw_device_type = av_hwdevice_find_type_by_name(hw_name.c_str());
    }
    const AVCodec * target = FindCodecByCustomFunc([name, hw_device_type, this](const AVCodec * codec){

        if (codec && name == codec->name && (codec->decode != nullptr || codec->receive_frame != nullptr))
        {
            for (int index = 0; ; index++)
            {
                const AVCodecHWConfig * config = avcodec_get_hw_config(codec, index);
                if (!config)
                {
                    break;
                }
                // hardware no support
                if (config->device_type == AV_HWDEVICE_TYPE_CUDA)
                {
                    continue;
                }
                // if (hw_device_type != AV_HWDEVICE_TYPE_NONE && config->device_type != hw_device_type)
                // {
                //     continue;
                // }
                if (config->pix_fmt != AV_PIX_FMT_NONE && config->methods & AV_CODEC_HW_CONFIG_METHOD_HW_DEVICE_CTX && config->device_type == hw_device_type)
                {
                    m_codec_hw_config = config;
                    return true;
                }
            }
        }
        return false;

    });
    return target;
}

const AVCodec * DecoderContext::FindCodecByIdAndHw(AVCodecID codec_id)
{
    const AVCodec * target = FindCodecByCustomFunc([codec_id, this](const AVCodec * codec){

        if (codec && codec->id == codec_id && (codec->decode != nullptr || codec->receive_frame != nullptr))
        {
            for (int index = 0; ; index++)
            {
                const AVCodecHWConfig * config = avcodec_get_hw_config(codec, index);
                if (!config)
                {
                    break;
                }
                if (config->pix_fmt != AV_PIX_FMT_NONE && config->methods & AV_CODEC_HW_CONFIG_METHOD_HW_DEVICE_CTX)
                {
                    m_codec_hw_config = config;
                    return true;
                }
            }
        }
        return false;
    });
    return target;
}

const AVCodec * DecoderContext::FindCodecByIdAndHw(AVCodecID codec_id, AVHWDeviceType hw_device_type)
{
    const AVCodec * target = FindCodecByCustomFunc([codec_id, hw_device_type, this](const AVCodec * codec){

        if (codec && codec->id == codec_id && (codec->decode != nullptr || codec->receive_frame != nullptr))
        {
            if (hw_device_type == AV_HWDEVICE_TYPE_NONE)
            {
                return true;
            }
            for (int index = 0; ; index++)
            {
                const AVCodecHWConfig * config = avcodec_get_hw_config(codec, index);
                if (!config)
                {
                    break;
                }
                if (config->device_type == hw_device_type && config->pix_fmt != AV_PIX_FMT_NONE && config->methods & AV_CODEC_HW_CONFIG_METHOD_HW_DEVICE_CTX)
                {
                    m_codec_hw_config = config;
                    return true;
                }
            }
        }
        return false;
    });
    return target;
}

int DecoderContext::GetHwMethod()
{
    if (m_codec_hw_config)
    {
        return m_codec_hw_config->methods;
    }
    return -1;
}

bool DecoderContext::GetHwInitStatus()
{
    return m_hw_frame_init_status;
}

const AVCodec * DecoderContext::FindCodecByCustomFunc(std::function<bool (const AVCodec *)> &&func)
{
    const AVCodec * target = nullptr;

    void * index = 0;
    const AVCodec * ptr = nullptr;

    while ((ptr = av_codec_iterate(&index)))
    {
        if (!ptr)
        {
            continue;
        }
        bool result = func(ptr);
        if (result)
        {
            target = ptr;
            break;
        }
    }
    return target;

}

AVCodecContext * DecoderContext::AvCodecCtx()
{
    return m_decoder_ctx;
}


EncoderContext::EncoderContext()
{
    m_encoder_ctx = nullptr;
    m_hw_frame_init_status = false;
}

EncoderContext::~EncoderContext()
{
    if (m_encoder_ctx)
    {
        avcodec_close(m_encoder_ctx);
        avcodec_free_context(&m_encoder_ctx);
        m_encoder_ctx = nullptr;
        PrintLog("release codec ctx resouce");
    }
}

bool EncoderContext::GetHwInitStatus()
{
    return m_hw_frame_init_status;
}

bool EncoderContext::CreateEncoderContextByName(const std::string & name, const std::string & hw_name)
{
    const AVCodec * codec = nullptr;
    if (!hw_name.empty())
    {
        codec = FindCodecByNameAndHw(name, hw_name);
    }
    else
    {
        codec = FindCodecByName(name);
    }
    if (!codec)
    {
        return false;
    }
    m_encoder_ctx = avcodec_alloc_context3(codec);
    if (!m_encoder_ctx)
    {
        m_encoder_ctx = nullptr;
        return false;
    }
    // m_codec = codec;
    return true;
}

bool EncoderContext::ProcessHwConfig()
{
    if (m_codec_hw_config && !m_hw_frame_init_status)
    {
        if (m_codec_hw_config->methods & AV_CODEC_HW_CONFIG_METHOD_HW_DEVICE_CTX)
        {
            PrintLog("ProcessHwConfig AV_CODEC_HW_CONFIG_METHOD_HW_DEVICE_CTX");
            m_hw_frame_init_status = false;

        }
        else if (m_codec_hw_config->methods & AV_CODEC_HW_CONFIG_METHOD_HW_FRAMES_CTX)
        {
            PrintLog("ProcessHwConfig AV_CODEC_HW_CONFIG_METHOD_HW_FRAMES_CTX");
            m_hw_frame_init_status = false;
        }
    }
    else
    {
        m_hw_frame_init_status = true;
    }
    return true;
}

void EncoderContext::SetAvBufferRef(AVBufferRef *ref_buffer)
{

}

void EncoderContext::ConfigParam(const VideoParamConfig & param_config)
{
    if (m_encoder_ctx && !param_config.m_copy_decoder_param)
    {
        if (param_config.m_height > 0 && param_config.m_width > 0)
        {
            m_encoder_ctx->height  = param_config.m_height;
            m_encoder_ctx->width   = param_config.m_width;
        }
        if (param_config.m_frame)
        {
            m_encoder_ctx->color_range     = param_config.m_frame->color_range;
            m_encoder_ctx->color_primaries = param_config.m_frame->color_primaries;
            m_encoder_ctx->color_trc       = param_config.m_frame->color_trc;
            m_encoder_ctx->colorspace      = param_config.m_frame->colorspace;
            m_encoder_ctx->chroma_sample_location = param_config.m_frame->chroma_location;
            m_encoder_ctx->field_order     = AV_FIELD_PROGRESSIVE;
        }
        if (param_config.m_config_callback && m_encoder_ctx)
        {
            param_config.m_config_callback(m_encoder_ctx);
        }
    }
}

void EncoderContext::ConfigParam(const AudioParamConfig &param_config)
{
    if (m_encoder_ctx && !param_config.m_copy_decoder_param)
    {
        if (param_config.m_format >= 0)
        {
            m_encoder_ctx->sample_fmt = (AVSampleFormat)param_config.m_format;
        }
        if (param_config.m_channel > 0)
        {
            m_encoder_ctx->channels = param_config.m_channel;
        }
        if (param_config.m_channel_layout > 0)
        {
            m_encoder_ctx->channel_layout = param_config.m_channel_layout;
        }
        if (param_config.m_frame)
        {
            // copy from ffmpeg source code
        }
        if (param_config.m_config_callback)
        {
            param_config.m_config_callback(m_encoder_ctx);
        }
    }
}

bool EncoderContext::OpenCodecCtx()
{
    if (m_codec_hw_config && !ProcessHwConfig())
    {
        PrintLog("has problem");
        return false;
    }
    else if (!m_codec_hw_config)
    {
        m_hw_frame_init_status = true;
    }
    if (m_hw_frame_init_status)
    {
        AVDictionary * thread_opt = nullptr;
        av_dict_set(&thread_opt, "threads", "auto", 0);
        int result = avcodec_open2(m_encoder_ctx, m_encoder_ctx->codec, &thread_opt);
        if (result < 0)
        {
            av_dict_free(&thread_opt);
            PrintLog(result);
            return false;
        }
    }
    return true;
}

bool EncoderContext::CreateEncoderContextById(AVCodecID codec_id, const std::string & hw_name)
{
    const AVCodec * target = FindCodecByIdAndHw(codec_id, hw_name);
    if (target)
    {
        m_encoder_ctx = avcodec_alloc_context3(target);
        return true;
    }
    return false;
}

const AVCodec * EncoderContext::FindCodecByName(const std::string & name)
{
    const AVCodec * target = FindCodecByCusttomFunc([name](const AVCodec * codec){

        if (codec && codec->name == name && (codec->encode2 || codec->receive_packet || codec->encode_sub))
        {
            return true;
        }
        return false;
    });
    return target;
}

const AVCodec * EncoderContext::FindCodecById(AVCodecID codec_id)
{
    const AVCodec * target = FindCodecByCusttomFunc([codec_id](const AVCodec * codec){

        if (codec && codec->id == codec_id && (codec->encode2 || codec->receive_packet || codec->encode_sub))
        {
            return true;
        }
        return false;
    });
    return target;
}

const AVCodec * EncoderContext::FindCodecByNameAndHw(const std::string & name, const std::string & hw_name)
{
    AVHWDeviceType hw_device_type = AV_HWDEVICE_TYPE_NONE;
    if (!hw_name.empty())
    {
        hw_device_type = av_hwdevice_find_type_by_name(hw_name.c_str());
    }
    auto * target = FindCodecByCusttomFunc([name,hw_device_type,this](const AVCodec * codec){

        if (codec && codec->name == name && (codec->encode2 || codec->receive_packet || codec->encode_sub))
        {
            if (hw_device_type != AV_HWDEVICE_TYPE_NONE)
            {
                for (int index = 0; ; index++)
                {
                    const AVCodecHWConfig * config = avcodec_get_hw_config(codec, index);
                    if (!config)
                    {
                        break;
                    }
                    if (config->device_type == hw_device_type && config->pix_fmt != AV_PIX_FMT_NONE)
                    {
                        m_codec_hw_config = config;
                        return true;
                    }
                }
            }
            else
            {
                return true;
            }
        }
        return false;
    });
    return target;
}

const AVCodec * EncoderContext::FindCodecByIdAndHw(AVCodecID codec_id, const std::string & hw_name)
{
    AVHWDeviceType hw_device_type = AV_HWDEVICE_TYPE_NONE;
    if (!hw_name.empty())
    {
        hw_device_type = av_hwdevice_find_type_by_name(hw_name.c_str());
    }
    auto * target = FindCodecByCusttomFunc([codec_id,hw_device_type,this](const AVCodec * codec){

        if (codec && codec->id == codec_id && (codec->encode2 || codec->receive_packet || codec->encode_sub))
        {
            if (AV_HWDEVICE_TYPE_NONE != hw_device_type)
            {
                for (int index = 0; ; index++)
                {
                    const AVCodecHWConfig * config = avcodec_get_hw_config(codec, index);
                    if (!config)
                    {
                        break;
                    }
                    if (config->device_type == hw_device_type && config->pix_fmt != AV_PIX_FMT_NONE)
                    {
                        m_codec_hw_config = config;
                        return true;
                    }
                }
            }
            else
            {
                return true;
            }
        }
        return false;
    });
    return target;
}

const AVCodec * EncoderContext::FindCodecByCusttomFunc(std::function<bool(const AVCodec * ptr)> && func)
{
    const AVCodec * target = nullptr;

    void * index = 0;
    const AVCodec * ptr = nullptr;

    while ((ptr = av_codec_iterate(&index)))
    {
        if (!ptr)
        {
            continue;
        }
        bool result = func(ptr);
        if (result)
        {
            target = ptr;
            break;
        }
    }
    return target;
}

AVCodecContext * EncoderContext::AvCodecCtx()
{
    return m_encoder_ctx;
}
